package com.zy.je.common.core.utils.thread;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author zy
 * 扩展ScheduledThreadPoolExecutor，使得可以控制队列大小
 */
@Slf4j
public class BoundedScheduledExecutor {

    private final ScheduledThreadPoolExecutor scheduler;
    private final int queueCapacity;
    private final RejectedExecutionHandler rejectedExecutionHandler;

    private final static AtomicInteger THREAD_NUM = new AtomicInteger(1);

    public BoundedScheduledExecutor(int corePoolSize, int queueCapacity) {
        this(corePoolSize, queueCapacity, r -> {
                    Thread t = new Thread(r);
                    long id = Thread.currentThread().getId();
                    t.setName("bounded-scheduled-pool-" + THREAD_NUM.getAndIncrement() + "-" + id);
                    return t;
                }
        );
    }

    public BoundedScheduledExecutor(int corePoolSize, int queueCapacity, ThreadFactory threadFactory) {
        this.scheduler = new ScheduledThreadPoolExecutor(corePoolSize, threadFactory);
        this.queueCapacity = queueCapacity;
        // 设置默认的拒绝策略为丢弃任务
        this.rejectedExecutionHandler = new ThreadPoolExecutor.DiscardPolicy();
    }

    // 提供调度方法
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command, long initialDelay, long delay, TimeUnit unit) {
        return scheduler.scheduleWithFixedDelay(() -> {
            log.debug("queue size {}",scheduler.getQueue().size());
            if (scheduler.getQueue().size() >= queueCapacity) {
                // 如果队列已满，调用拒绝策略
                log.info("rejected execution....");
                rejectedExecutionHandler.rejectedExecution(command, scheduler);
                return;
            }
            // 如果队列未满，则提交任务
            scheduler.submit(command);
        }, initialDelay, delay, unit);
    }

    public ScheduledFuture<?> scheduleAtFixedRate(Runnable command, long initialDelay, long delay, TimeUnit unit) {
        return scheduler.scheduleAtFixedRate(() -> {
            log.debug("queue size {}",scheduler.getQueue().size());
            if (scheduler.getQueue().size() >= queueCapacity) {
                // 如果队列已满，调用拒绝策略
                log.info("rejected execution....");
                rejectedExecutionHandler.rejectedExecution(command, scheduler);
                return;
            }
            // 如果队列未满，则提交任务
            scheduler.submit(command);
        }, initialDelay, delay, unit);
    }

    // 关闭线程池
    public void shutdown() {
        scheduler.shutdown();
    }

    // 确保所有任务完成后再关闭
    public void shutdownAndAwaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        scheduler.shutdown();
        if (!scheduler.awaitTermination(timeout, unit)) {
            scheduler.shutdownNow();
            if (!scheduler.awaitTermination(timeout, unit)) {
                log.error("Scheduler did not terminate");
            }
        }
    }

    public static void main(String[] args) {
        BoundedScheduledExecutor boundedExecutor = new BoundedScheduledExecutor(1, 1);

        // 使用scheduleWithFixedDelay方法安排任务
        boundedExecutor.scheduleAtFixedRate(() -> {
            try {
                Thread.sleep(2 * 1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            log.info("Executing task at " + System.currentTimeMillis());
        }, 0, 1, TimeUnit.SECONDS);

        // 记得在适当的时候关闭线程池
        // boundedExecutor.shutdown();
    }
}